{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c1f4145e-66fc-45da-9955-04db80ea206a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "from shapely.affinity import rotate, translate\n",
    "from shapely.geometry import LineString, Polygon\n",
    "from shapely.ops import unary_union\n",
    "\n",
    "from prime_dataset import FloorplanDataset, floorplan_collate\n",
    "from visualize import get_hard_color, visualize_prime\n",
    "\n",
    "\n",
    "from utils import normalize_polygon, polygons_have_same_shape, check_fixed_const, check_preplaced_const, check_mib_const, check_boundary_const, check_clust_const\n",
    "from cost import calculate_weighted_b2b_wirelength, calculate_weighted_p2b_wirelength\n",
    "\n",
    "##sample input for input solution tensor with 21-partitions\n",
    "\"\"\"\n",
    "[tensor([[180.,  45.],\n",
    "         [180.,  60.],\n",
    "         [186.,  60.],\n",
    "         [186.,  45.],\n",
    "         [180.,  45.]], dtype=torch.float16),\n",
    " tensor([[180.,  75.],\n",
    "         [180., 105.],\n",
    "         [186., 105.],\n",
    "         [186.,  75.],\n",
    "         [180.,  75.]], dtype=torch.float16),\n",
    " tensor([[150.,  30.],\n",
    "         [150.,   0.],\n",
    "         [120.,   0.],\n",
    "         [120.,  30.],\n",
    "         [150.,  30.]], dtype=torch.float16),\n",
    " tensor([[135.,  60.],\n",
    "         [135.,  75.],\n",
    "         [186.,  75.],\n",
    "         [186.,  60.],\n",
    "         [135.,  60.]], dtype=torch.float16),\n",
    " tensor([[ 60., 135.],\n",
    "         [ 90., 135.],\n",
    "         [ 90.,  75.],\n",
    "         [ 60.,  75.],\n",
    "         [ 60., 135.]], dtype=torch.float16),\n",
    " tensor([[120.,  60.],\n",
    "         [120.,  75.],\n",
    "         [135.,  75.],\n",
    "         [135.,  45.],\n",
    "         [120.,  45.],\n",
    "         [120.,  60.]], dtype=torch.float16),\n",
    " tensor([[90., 60.],\n",
    "         [90.,  0.],\n",
    "         [60.,  0.],\n",
    "         [60., 60.],\n",
    "         [90., 60.]], dtype=torch.float16),\n",
    " tensor([[120.,  30.],\n",
    "         [120.,   0.],\n",
    "         [ 90.,   0.],\n",
    "         [ 90.,  45.],\n",
    "         [135.,  45.],\n",
    "         [135.,  60.],\n",
    "         [150.,  60.],\n",
    "         [150.,  30.],\n",
    "         [120.,  30.]], dtype=torch.float16),\n",
    " tensor([[186.,  45.],\n",
    "         [186.,  30.],\n",
    "         [150.,  30.],\n",
    "         [150.,  60.],\n",
    "         [180.,  60.],\n",
    "         [180.,  45.],\n",
    "         [186.,  45.]], dtype=torch.float16),\n",
    " tensor([[180.,  75.],\n",
    "         [120.,  75.],\n",
    "         [120., 105.],\n",
    "         [180., 105.],\n",
    "         [180.,  75.]], dtype=torch.float16),\n",
    " tensor([[135., 165.],\n",
    "         [135., 135.],\n",
    "         [150., 135.],\n",
    "         [150., 105.],\n",
    "         [120., 105.],\n",
    "         [120., 165.],\n",
    "         [135., 165.]], dtype=torch.float16),\n",
    " tensor([[  0., 186.],\n",
    "         [ 30., 186.],\n",
    "         [ 30., 135.],\n",
    "         [  0., 135.],\n",
    "         [  0., 186.]], dtype=torch.float16),\n",
    " tensor([[150., 186.],\n",
    "         [150., 135.],\n",
    "         [135., 135.],\n",
    "         [135., 165.],\n",
    "         [120., 165.],\n",
    "         [120., 186.],\n",
    "         [150., 186.]], dtype=torch.float16),\n",
    " tensor([[150.,  30.],\n",
    "         [186.,  30.],\n",
    "         [186.,   0.],\n",
    "         [150.,   0.],\n",
    "         [150.,  30.]], dtype=torch.float16),\n",
    " tensor([[ 60.,  75.],\n",
    "         [ 30.,  75.],\n",
    "         [ 30.,  60.],\n",
    "         [  0.,  60.],\n",
    "         [  0., 105.],\n",
    "         [ 60., 105.],\n",
    "         [ 60.,  75.]], dtype=torch.float16),\n",
    " tensor([[60.,  0.],\n",
    "         [ 0.,  0.],\n",
    "         [ 0., 30.],\n",
    "         [60., 30.],\n",
    "         [60.,  0.]], dtype=torch.float16),\n",
    " tensor([[45., 45.],\n",
    "         [60., 45.],\n",
    "         [60., 30.],\n",
    "         [ 0., 30.],\n",
    "         [ 0., 60.],\n",
    "         [30., 60.],\n",
    "         [30., 75.],\n",
    "         [45., 75.],\n",
    "         [45., 45.]], dtype=torch.float16),\n",
    " tensor([[60., 75.],\n",
    "         [90., 75.],\n",
    "         [90., 60.],\n",
    "         [60., 60.],\n",
    "         [60., 45.],\n",
    "         [45., 45.],\n",
    "         [45., 75.],\n",
    "         [60., 75.]], dtype=torch.float16),\n",
    " tensor([[ 45., 105.],\n",
    "         [  0., 105.],\n",
    "         [  0., 135.],\n",
    "         [ 30., 135.],\n",
    "         [ 30., 186.],\n",
    "         [ 45., 186.],\n",
    "         [ 45., 105.]], dtype=torch.float16),\n",
    " tensor([[180., 186.],\n",
    "         [186., 186.],\n",
    "         [186., 105.],\n",
    "         [150., 105.],\n",
    "         [150., 186.],\n",
    "         [180., 186.]], dtype=torch.float16),\n",
    " tensor([[ 45., 186.],\n",
    "         [120., 186.],\n",
    "         [120.,  45.],\n",
    "         [ 90.,  45.],\n",
    "         [ 90., 135.],\n",
    "         [ 60., 135.],\n",
    "         [ 60., 105.],\n",
    "         [ 45., 105.],\n",
    "         [ 45., 186.]], dtype=torch.float16)] \n",
    "\"\"\"\n",
    "\n",
    "\n",
    "\n",
    "def estimate_cost(bdata: list, layout_index: int):\n",
    "    \"\"\"\n",
    "    Estimate the cost of a layout by evaluating area and wirelength violations.\n",
    "\n",
    "    Args:\n",
    "        bdata (list): The block data containing polygon coordinates.\n",
    "        layout_index (int): The index of the layout to evaluate.\n",
    "    \"\"\"\n",
    "    # Read baseline design\n",
    "    root = './'\n",
    "    ds = FloorplanDataset(root)\n",
    "    target_data = ds.__getitem__(layout_index)\n",
    "\n",
    "    target_area_budgets = target_data['input'][0]\n",
    "    target_b2b_edges = target_data['input'][1]\n",
    "    target_p2b_edges = target_data['input'][2]\n",
    "    target_pins_pos = target_data['input'][3]\n",
    "    target_constraints = target_data['input'][4]\n",
    "    target_metrics = target_data['label'][1]\n",
    "    target_poly = target_data['label'][0]\n",
    "\n",
    "    target_b2b_wl = target_metrics[-2].item()\n",
    "    target_p2b_wl = target_metrics[-1].item()\n",
    "    target_layout_area = target_metrics[0].item()\n",
    "\n",
    "    if len(bdata) != len(target_poly):\n",
    "        print('ERROR: incorrect number of polygons in the solution')\n",
    "        exit()\n",
    "\n",
    "    # Estimate area and validate target area budgets\n",
    "    target_sol = [Polygon(targ_poly.tolist()) for targ_poly in target_poly]\n",
    "\n",
    "    sol_area_budgets = torch.zeros_like(target_area_budgets)\n",
    "    fp_sol = []\n",
    "    W, H = 0, 0\n",
    "    min_x, min_y, max_x, max_y = float('inf'), float('inf'), 0, 0\n",
    "\n",
    "    for ind, elem in enumerate(bdata):\n",
    "        poly_elem = Polygon(elem.tolist())\n",
    "        sol_area_budgets[ind] = poly_elem.area\n",
    "        fp_sol.append(poly_elem)\n",
    "\n",
    "        min_x, min_y, max_x, max_y = poly_elem.bounds\n",
    "        W = max(W, max_x)\n",
    "        H = max(H, max_y)\n",
    "\n",
    "    sol_area_cost = float(W * H)\n",
    "    delta_budgets = sol_area_budgets - target_area_budgets\n",
    "    area_viol = torch.nonzero(delta_budgets < 0).squeeze()\n",
    "\n",
    "    # Estimate wirelength\n",
    "    centroids = torch.tensor(\n",
    "        [list(Polygon(fp_sol[i]).centroid.coords[0]) for i in range(len(fp_sol))],\n",
    "        dtype=torch.float32,\n",
    "    )\n",
    "\n",
    "    sol_b2b_wl = calculate_weighted_manhattan_distance(\n",
    "        centroids, target_b2b_edges\n",
    "    )\n",
    "\n",
    "    sol_p2b_wl = calculate_pin_wirelength(\n",
    "        centroids, target_p2b_edges, target_pins_pos\n",
    "    )\n",
    "\n",
    "\n",
    "\n",
    "    # Estimate constraint cost\n",
    "    fixed_const = target_constraints[:, 0]\n",
    "    preplaced_const = target_constraints[:, 1]\n",
    "    mib_const = target_constraints[:, 2]\n",
    "    clust_const = target_constraints[:, 3]\n",
    "    bound_const = target_constraints[:, 4]\n",
    "\n",
    "    # Prepare the output dictionary\n",
    "    results = {\n",
    "        'placement_constraints': {\n",
    "            'fixed': 0,\n",
    "            'preplaced': 0,\n",
    "            'mib': 0,\n",
    "            'cluster': 0,\n",
    "            'boundary': 0\n",
    "        },\n",
    "        'wl_difference': {\n",
    "            'b2b': sol_b2b_wl - target_b2b_wl,\n",
    "            'p2b': sol_p2b_wl - target_p2b_wl\n",
    "        },\n",
    "        'layout_area_difference': sol_area_cost - target_layout_area,\n",
    "        'partition_indices_with_area_violations': area_viol.tolist()\n",
    "    }\n",
    "\n",
    "    results['placement_constraints']['fixed'] = check_fixed_const(\n",
    "        torch.nonzero(fixed_const), fp_sol, target_sol\n",
    "    )\n",
    "    results['placement_constraints']['preplaced'] = check_preplaced_const(\n",
    "        torch.nonzero(preplaced_const), fp_sol, target_sol\n",
    "    )\n",
    "    results['placement_constraints']['mib'] = check_mib_const(\n",
    "        mib_const, fp_sol, target_sol\n",
    "    )\n",
    "    results['placement_constraints']['cluster'] = check_clust_const(\n",
    "        clust_const, fp_sol, target_sol\n",
    "    )\n",
    "    results['placement_constraints']['boundary'] = check_boundary_const(\n",
    "        bound_const, fp_sol, target_sol, W, H\n",
    "    )\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e57af271-92c9-4be9-94ee-4d24f3fd3659",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1102be9b-6dfb-42a6-b914-4bddc92b0add",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
